use super::sparse;
use crate::utils::global2local;
use ndarray::prelude::*;
use rand::prelude::*;
use std::collections::BTreeSet;
pub struct HyperGraph<A: Clone + std::marker::Send + std::marker::Sync> {
    node_incidence: sparse::CSR<A>,
    edge_incidence: sparse::CSR<A>,
}

impl<A> HyperGraph<A>
where
    A: Clone + std::marker::Send + std::marker::Sync,
{
    pub fn new(
        data: ArrayView1<A>,
        row_indices: ArrayView1<usize>,
        col_indices: ArrayView1<usize>,
        num_nodes: usize,
        num_edges: usize,
    ) -> HyperGraph<A> {
        let ni = sparse::COO::<A>::new(
            &data.view(),
            &row_indices.view(),
            &col_indices.view(),
            num_nodes,
            num_edges,
        )
        .tocsr();
        // let ei = sparse::COO::<A>::new(&data, &col_indices, &row_indices, num_edges, num_nodes).tocsr();
        let ei = ni.transpose();
        HyperGraph::<A> {
            node_incidence: ni,
            edge_incidence: ei,
        }
    }

    pub fn num_node(&self) -> usize {
        self.node_incidence.shape().0
    }

    pub fn num_edge(&self) -> usize {
        self.edge_incidence.shape().0
    }

    pub fn restart_random_walk(
        &self,
        nodes: &ArrayView1<usize>,
        length: usize,
        num_path: usize,
        rate: f32,
    ) -> Array3<usize> {
        let mut result = Array3::<usize>::zeros((nodes.len(), num_path, length));
        for i in 0..nodes.len() {
            let mut rng = rand::thread_rng();
            for p in 0..num_path {
                let mut next_node_ind = nodes[i];
                for j in 0..length {
                    let (col_ind, _) = self.node_incidence.row(next_node_ind);
                    if col_ind.len() == 0 {
                        // result[[i, p, j]] = self.num_node();
                        next_node_ind = nodes[i];
                        continue;
                    }
                    let (temp, _) = self
                        .edge_incidence
                        .row(col_ind[rng.gen_range(0..col_ind.len())]);
                    next_node_ind = temp[rng.gen_range(0..temp.len())];
                    result[[i, p, j]] = next_node_ind;
                    if rng.gen::<f32>() < rate {
                        next_node_ind = nodes[i];
                    }
                }
            }
        }
        result
    }

    pub fn sample_linkprediction_by_restart_random_walk(
        &self,
        nodes: &ArrayView1<usize>,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
        rate: f32,
        num_e: usize,
        num_n: usize,
    ) -> (
        Array1<usize>,
        Array2<usize>,
        Array3<usize>,
        Array3<usize>,
        Array3<usize>,
    ) {
        let results = self.restart_random_walk(nodes, length, num_path, rwr_rate);
        let new_nodes = Array1::from_iter(BTreeSet::from_iter(
            results.iter().chain(nodes.iter()).cloned(),
        ));
        let mut temp = BTreeSet::new();
        for &n in new_nodes.iter() {
            temp.extend(self.node_incidence.row(n).0.iter());
        }
        let new_edges = Array1::from_iter(temp);
        let nind = self.node_incidence.submatrix(&new_nodes, &new_edges);
        let eind = nind.transpose();
        let local_nodes = global2local(&new_nodes.view(), nodes);
        let (a, b, c, d) =
            sample_linkprediction(&nind, &eind, &local_nodes.view(), num_e, num_n, rate);
        (new_nodes, a, b, c, d)
    }

    pub fn sample_by_restart_random_walk(
        &self,
        nodes: &ArrayView1<usize>,
        length: usize,
        num_path: usize,
        rwr_rate: f32,
        num_n: usize,
    ) -> (Array1<usize>, Array2<usize>) {
        let results = self.restart_random_walk(nodes, length, num_path, rwr_rate);
        let new_nodes = Array1::from_iter(BTreeSet::from_iter(
            results.iter().chain(nodes.iter()).cloned(),
        ));
        let mut temp = BTreeSet::new();
        for &n in new_nodes.iter() {
            temp.extend(self.node_incidence.row(n).0.iter());
        }
        let new_edges = Array1::from_iter(temp);
        let nind = self.node_incidence.submatrix(&new_nodes, &new_edges);
        let eind = nind.transpose();
        // let local_nodes = global2local(&new_nodes.view(), nodes);

        let mut rng = rand::thread_rng();
        let mut iind = Vec::<Array1<usize>>::new();
        for i in 0..eind.shape().0 {
            let entemp = eind.row(i).0.to_vec();
            if entemp.len() > num_n {
                iind.push(
                    entemp
                        .choose_multiple(&mut rng, num_n)
                        .cloned()
                        .collect::<Array1<usize>>(),
                );
            } else {
                iind.push(
                    (0..num_n)
                        .map(|_| entemp[rng.gen_range(0..entemp.len())])
                        .collect::<Array1<usize>>(),
                );
            }
        }
        let mut eind_sample = Array2::zeros((iind.len(), num_n));
        for (i, e) in iind.iter().enumerate() {
            eind_sample.slice_mut(s![i, ..]).assign(&e);
        }
        (new_nodes, eind_sample)
    }
}

fn sample_linkprediction<A: Clone + std::marker::Send + std::marker::Sync>(
    nind: &sparse::CSR<A>,
    eind: &sparse::CSR<A>,
    nodes: &ArrayView1<usize>,
    num_e: usize,
    num_n: usize,
    rate: f32,
) -> (Array2<usize>, Array3<usize>, Array3<usize>, Array3<usize>) {
    let mut rng = rand::thread_rng();
    let mut iind = Vec::<Array1<usize>>::new();
    let mut cursor = Array3::zeros((nodes.shape()[0], num_e, num_n));
    let mut psample = Array3::zeros((nodes.shape()[0], num_e, num_n));
    let mut nsample = Array3::zeros((nodes.shape()[0], num_e, num_n));

    for i in 0..eind.shape().0 {
        let entemp = eind.row(i).0.to_vec();
        if entemp.len() == 0 {
            continue;
        }
        if !nodes.iter().any(|x| entemp.contains(x)) {
            if entemp.len() > num_n {
                iind.push(
                    entemp
                        .choose_multiple(&mut rng, num_n)
                        .cloned()
                        .collect::<Array1<usize>>(),
                );
            } else {
                iind.push(
                    (0..num_n)
                        .map(|_| entemp[rng.gen_range(0..entemp.len())])
                        .collect::<Array1<usize>>(),
                );
            }
        }
    }
    for (i, &n) in nodes.iter().enumerate() {
        let (neigh, _) = nind.row(n);
        for j in 0..num_e {
            if neigh.len() == 0 {
                nsample.slice_mut(s![i, j, ..]).assign(&Array1::from_vec(
                    (0..nind.shape().0).choose_multiple(&mut rng, num_n),
                ));
                continue;
            }
            let etemp = neigh[rng.gen_range(0..neigh.len())]; // 边的ind
            let entemp = eind.row(etemp).0.to_vec(); // 边的组成，即节点
            let enntemp = if entemp.len() > num_n {
                // 边中节点数量与需要采样数量对比
                rng.gen_range(0..num_n) // 随机一个负样本数量
            } else {
                // 边中节点数量比需要采样的少
                rng.gen_range(0..entemp.len())
            };
            nsample.slice_mut(s![i, j, ..enntemp]).assign(
                &entemp
                    .choose_multiple(&mut rng, enntemp)
                    .cloned()
                    .collect::<Array1<usize>>(),
            );
            let mut nnntemp = (0..nind.shape().0).collect::<Vec<usize>>();
            nnntemp.retain(|n| !entemp.contains(n));
            nsample.slice_mut(s![i, j, enntemp..]).assign(
                &nnntemp
                    .choose_multiple(&mut rng, num_n - enntemp)
                    .cloned()
                    .collect::<Array1<usize>>(),
            );
        }
        match neigh.len() {
            0 => {
                psample.slice_mut(s![i, .., ..]).fill(nind.shape().0);
                cursor.slice_mut(s![i, .., ..]).fill(nind.shape().0);
            } // 不应该存在这种情况
            _ => {
                // divide neighbors of node to two set? sure? nope should like negative
                let mut edgetemp = Vec::new();
                for &j in neigh.iter() {
                    let entemp = eind.row(j).0;
                    let itemp = if entemp.len() == 1 {
                        entemp.to_vec()
                    } else {
                        let isize_ = (entemp.len() as f32 * rate).floor() as usize;
                        entemp
                            .to_vec()
                            .choose_multiple(&mut rng, isize_)
                            .cloned()
                            .collect::<Vec<usize>>()
                    };
                    if itemp.len() > num_n {
                        iind.push(
                            itemp
                                .choose_multiple(&mut rng, num_n)
                                .cloned()
                                .collect::<Array1<usize>>(),
                        );
                    } else {
                        iind.push(
                            (0..num_n)
                                .map(|_| entemp[rng.gen_range(0..itemp.len())])
                                .collect::<Array1<usize>>(),
                        );
                    }
                    edgetemp.push(itemp);
                }
                for j in 0..num_e {
                    let idx = rng.gen_range(0..neigh.len());
                    let pentemp = eind.row(neigh[idx]).0.to_vec();
                    if pentemp.len() > num_n {
                        psample.slice_mut(s![i, j, ..]).assign(
                            &pentemp
                                .choose_multiple(&mut rng, num_n)
                                .cloned()
                                .collect::<Array1<usize>>(),
                        );
                    } else {
                        psample.slice_mut(s![i, j, ..]).assign(
                            &(0..num_n)
                                .map(|_| pentemp[rng.gen_range(0..pentemp.len())])
                                .collect::<Array1<usize>>(),
                        );
                    }
                    let centemp = &edgetemp[idx];
                    if centemp.len() > num_n {
                        cursor.slice_mut(s![i, j, ..]).assign(
                            &centemp
                                .choose_multiple(&mut rng, num_n)
                                .cloned()
                                .collect::<Array1<usize>>(),
                        );
                    } else {
                        cursor.slice_mut(s![i, j, ..]).assign(
                            &(0..num_n)
                                .map(|_| centemp[rng.gen_range(0..centemp.len())])
                                .collect::<Array1<usize>>(),
                        );
                    }
                }
            }
        }
    }

    let mut eind_sample = Array2::zeros((iind.len(), num_n));
    for (i, e) in iind.iter().enumerate() {
        eind_sample.slice_mut(s![i, ..]).assign(&e);
    }
    (eind_sample, cursor, psample, nsample)
}
#[cfg(test)]
mod tests {
    use super::*;
    use polars::prelude::*;
    use rayon::prelude::*;
    use std::error::Error;
    use std::fs::File;
    use std::path::Path;

    #[test]
    fn hypergraph_new_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(
            r"/media/marcfish/backup/dataset/dataset/network/gowalla/user_local_5.parquet",
        );
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .select(["usre_id", "location_id"])?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let hg: HyperGraph<f32> = HyperGraph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            10328,
            83223,
        );
        assert_eq!(hg.num_node(), 10328);
        assert_eq!(hg.num_edge(), 83223);
        Ok(())
    }
    #[test]
    fn restart_random_walk_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(
            r"/media/marcfish/backup/dataset/dataset/network/gowalla/user_local_5.parquet",
        );
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .select(["usre_id", "location_id"])?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let hg: HyperGraph<f32> = HyperGraph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            10328,
            83223,
        );
        let a = hg.restart_random_walk(&Array1::from_iter(0..1024).view(), 3, 5, 0.3);
        assert_eq!(a.shape()[0], 1024);
        assert_eq!(hg.num_node(), 10328);
        assert_eq!(hg.num_edge(), 83223);
        Ok(())
    }

    #[test]
    fn sample_linkprediction_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(
            r"/media/marcfish/backup/dataset/dataset/table/meetup/Chicago_rsvp_3.parquet",
        );
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .select(["user_id", "event_id"])?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let hg: HyperGraph<f32> = HyperGraph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            43778,
            50534,
        );
        let (a, b, c, d, e) = hg.sample_linkprediction_by_restart_random_walk(
            &Array1::from_iter(1024 * 9..10240).view(),
            5,
            5,
            0.3,
            0.7,
            5,
            5,
        );
        // let (a, b, c) = sample_linkprediction(&hg.node_incidence, &hg.edge_incidence, &Array1::from_iter(1024*9..10240), 5, 10, 0.7);
        assert_ne!(a.shape()[0], 1024);
        assert_ne!(b.shape()[0], 1024);
        assert_eq!(c.shape()[0], 1024);
        assert_eq!(d.shape()[0], 1024);
        assert_eq!(e.shape()[0], 1024);
        // let a = hg.restart_random_walk(&Array1::from_iter(0..1024), 3, 5, 0.3);
        // assert_eq!(a.shape()[0], 1024);
        // assert_eq!(hg.num_node(), 10328);
        // assert_eq!(hg.num_edge(), 83223);
        Ok(())
    }

    #[test]
    fn sample_by_restart_random_walk_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(
            r"/media/marcfish/backup/dataset/dataset/table/meetup/Chicago_rsvp_3.parquet",
        );
        let file = File::open(path)?;
        let edge = ParquetReader::new(file)
            .finish()?
            .select(["user_id", "event_id"])?
            .to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let hg: HyperGraph<f32> = HyperGraph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            43778,
            50534,
        );
        let (a, b) = hg.sample_by_restart_random_walk(
            &Array1::from_iter(1024 * 9..10240).view(),
            5,
            5,
            0.3,
            10,
        );
        // let (a,b,c,d, e) = hg.sample_linkprediction_by_restart_random_walk(&Array1::from_iter(1024*9..10240).view(), 5, 5, 0.3, 0.7, 5, 5);
        // let (a, b, c) = sample_linkprediction(&hg.node_incidence, &hg.edge_incidence, &Array1::from_iter(1024*9..10240), 5, 10, 0.7);
        assert_ne!(a.shape()[0], 1024);
        assert_ne!(b.shape()[0], 1024);
        // assert_eq!(c.shape()[0], 1024);
        // assert_eq!(d.shape()[0], 1024);
        // assert_eq!(e.shape()[0], 1024);
        // let a = hg.restart_random_walk(&Array1::from_iter(0..1024), 3, 5, 0.3);
        // assert_eq!(a.shape()[0], 1024);
        // assert_eq!(hg.num_node(), 10328);
        // assert_eq!(hg.num_edge(), 83223);
        Ok(())
    }
    #[test]
    fn parallel_sample_linkprediction_works() -> Result<(), Box<dyn Error>> {
        let path = Path::new(
            r"/media/marcfish/backup/dataset/dataset/table/meetup/Chicago_rsvp_3.parquet",
        );
        let file = File::open(path)?;
        let df = ParquetReader::new(file).finish()?.sort(["created"], Default::default())?.select(["event_id", "user_id"])?;
        let edge = df.slice(0, (df.height() as f32 * 0.7).round() as usize).to_ndarray::<UInt32Type>(IndexOrder::C)?;
        let row = edge
            .select(Axis(1), &[0])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let col = edge
            .select(Axis(1), &[1])
            .remove_axis(Axis(1))
            .mapv(|x| x as usize);
        let hg: HyperGraph<f32> = HyperGraph::new(
            Array1::ones(edge.shape()[0]).view(),
            row.view(),
            col.view(),
            50534,
            43778,
        );
        let o: Vec<(
            Array1<usize>,
            Array2<usize>,
            Array3<usize>,
            Array3<usize>,
            Array3<usize>,
        )> = (0..9)
            .into_par_iter()
            .map(|x| {
                hg.sample_linkprediction_by_restart_random_walk(
                    &Array1::from_iter(x * 1024..(x + 1) * 1024).view(),
                    10,
                    5,
                    0.3,
                    0.7,
                    5,
                    5,
                )
            })
            .collect();
        // let (a,b,c,d, e) = hg.sample_linkprediction_by_restart_random_walk(&Array1::from_iter(0..1024), 5, 5, 0.3, 0.7, 5, 5);
        Ok(())
    }
}
